package com.gitee.cui.zookeeper.distributelock;

import com.gitee.cui.zk.ZKclient;
import lombok.extern.slf4j.Slf4j;
import org.apache.curator.framework.CuratorFramework;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 通过zk实现分布式锁，可重入公平锁
 * @author cuiqingqiang
 * @version 1.0
 * @date 2021/11/7 10:06
 **/
@Slf4j
public class ZKLock implements Lock {

    /**
     * 加锁节点
     */
    private static final String LOCK_PATH = "/test/lock";

    /**
     * 加锁临时子节点前缀
     */
    private static final String LOCK_PATH_PREFIX = LOCK_PATH + "/";

    /**
     * countdownlatch超时等待时间
     */
    private static final long WAIT_TIME = 1000;

    /**
     * 当前线程标识
     */
    private Thread current;

    /**
     * 可重入锁计数
     */
    private AtomicInteger count = new AtomicInteger(0);

    /**
     * 创建的自增临时节点路径
     */
    private String lockPath;

    /**
     * 创建的临时节点最后一部分
     */
    private String shortLockPath;

    /**
     * 前一个节点路径
     */
    private String previousNodePath;

    /**
     * Curator客户端
     */
    private CuratorFramework client;

    public ZKLock() {
        ZKclient instance = ZKclient.instance;
        instance.init();
        boolean nodeExist = instance.isNodeExist(LOCK_PATH);
        if (!nodeExist) {
            instance.createNode(LOCK_PATH, null);
        }
        client = instance.getClient();
    }

    @Override
    public boolean lock() {
        // 可重入锁判断，如果count=0，则表示第一次加锁，如果不等于0，判断是否是当前线程进入，如果是当前线程，计数加1
        synchronized (this) {
            if (count.get() == 0) {
                current = Thread.currentThread();
                count.incrementAndGet();
            } else {
                if (!current.equals(Thread.currentThread())) {
                    return false;
                }
                count.incrementAndGet();
                return true;
            }
        }
        try {
            // 尝试加锁
            boolean lock = tryLock();
            if (lock) {
                return true;
            }

            // 如果加锁不成功，则等待被唤醒进行下一次加锁
            while (!lock) {
                // 等待
                await();

                // 获取等待的列表
                List<String> waiters = getWaiters();
                if (checkLocked(waiters)) {
                    lock = true;
                }
            }
            return true;
        } catch (Exception e) {
            count.decrementAndGet();
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 检查是否获取到锁
     * @param waiters 等待列表
     * @return 是否获取到锁
     */
    private boolean checkLocked(List<String> waiters) {
        // 检查当前节点，是否在等待列表中顺序是最小的，如果是最小的，则获取到锁
        Collections.sort(waiters);

        if (shortLockPath.equals(waiters.get(0))) {
            log.info("成功获取分布式锁，节点为: {}", shortLockPath);
            return true;
        }
        return false;
    }

    /**
     * 等待被唤醒并尝试获取锁
     */
    private void await() throws Exception {
        // 监听前一个节点的删除事件
        if (previousNodePath == null) {
            throw new Exception("previousNodePath is null");
        }

        CountDownLatch countDownLatch = new CountDownLatch(1);

        Watcher watcher = new Watcher() {
            @Override
            public void process(WatchedEvent watchedEvent) {
                log.info("前一个节点的删除事件触发：{}", watchedEvent);
                // 唤醒线程继续执行
                countDownLatch.countDown();
            }
        };

        // 添加监听事件
        client.getData().usingWatcher(watcher).forPath(previousNodePath);

        // 等待被唤醒或者超时
        countDownLatch.await(WAIT_TIME, TimeUnit.MILLISECONDS);
    }

    /**
     * 获取所有等待列表，序号列表
     * @return 序号列表
     */
    private List<String> getWaiters() {
        List<String> waiters = new ArrayList<>();
        try {
            waiters = client.getChildren()
                    .forPath(LOCK_PATH);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return waiters;
    }

    /**
     * 尝试加锁
     * @return 加锁是否成功
     */
    private boolean tryLock() throws Exception {
        // 创建临时顺序子节点，判断自己的序号是不是最小的
        ZKclient instance = ZKclient.instance;
        lockPath = instance.createEphemeralSeqNode(LOCK_PATH_PREFIX);
        if (lockPath == null) {
            throw new Exception("zk error");
        }

        // 判断自己的序号是不是所有节点中最小的那个
        // 先获取节点的序号部分
        shortLockPath = getShortPath(lockPath);
        // 获取所有节点
        List<String> waiters = getWaiters();

        if (checkLocked(waiters)) {
            return true;
        }

        // 判断自己在所有等待列表中的位置
        int index = Collections.binarySearch(waiters, shortLockPath);
        if (index < 0) {
            // 由于多方面原因，没有在等待列表中找到自己，则可能节点被删除了
            throw new Exception("节点没有找到: " + shortLockPath);
        }

        // 找到自己前一个节点，并监听前一个节点的删除事件
        previousNodePath = LOCK_PATH_PREFIX + waiters.get(index - 1);
        return false;
    }

    /**
     * 从创建的节点中获取序号的部分
     * @param lockPath 创建的节点路径
     * @return 序号部分
     */
    private String getShortPath(String lockPath) {
        String shortLockPath = null;
        int index = lockPath.lastIndexOf(LOCK_PATH_PREFIX);
        if (index >= 0) {
            index += LOCK_PATH_PREFIX.length();
            shortLockPath = index <= lockPath.length() ? lockPath.substring(index) : null;
        }
        return shortLockPath;
    }

    @Override
    public boolean unlock() {
        if (!current.equals(Thread.currentThread())) {
            return false;
        }

        int count = this.count.decrementAndGet();
        if (count < 0) {
            throw new IllegalMonitorStateException("计数错误: " + count);
        }

        // 说明重入锁还存在锁
        if (count != 0) {
            return true;
        }

        // 释放锁，就是删除该临时节点
        ZKclient instance = ZKclient.instance;
        boolean nodeExist = instance.isNodeExist(lockPath);
        if (nodeExist) {
            instance.deleteNode(lockPath);
        }
        return false;
    }
}
